# pylint: disable=unused-wildcard-import
# pylint: disable=wildcard-import,invalid-name,missing-class-docstring,missing-function-docstring
import taichi as ti
from .type_enum import MaterialType
from .gtaichi import *

@ti.func
def refraction(wi,n,mu):
    """
    折射
    wi入射光,n法线,mu相对折射率
    """
    wo = ti.Vector([0.,0.,0.])
    cos_theta_0 = ti.math.dot(wi,n)
    if cos_theta_0<0:
        mu = 1/mu
        cos_theta_0 = -cos_theta_0
        n = -n
    sin_theta_0 = ti.math.sqrt(1.-cos_theta_0*cos_theta_0)
    sin_theta_1 = sin_theta_0/mu
    cos_theta_1 = ti.math.sqrt(1.-sin_theta_1*sin_theta_1)
    wo = ti.math.normalize(-wi/mu+(cos_theta_0/mu-cos_theta_1)*n)
    return wo,-n

@ti.func
def bxdf(wo:vec3,ints:Intersection,m:Material,c):
    """
    表面反射特性权重
    """
    n = ints.n
    wi = ti.Vector([0.,0.,0.])
    weight = ti.Vector([1.,1.,1.])
    if m.type==MaterialType.BRDF.value:
        micro_n = importance_sample_ggx(random_vec2(),
                                        n,m.a*m.a)
        wi = ti.math.reflect(-wo,micro_n)
        dot_wi_n = ti.max(ti.math.dot(wi,n),0.) #<0将导致dot_wi_n=0
        pdf = importance_sample_ggx_pdf(wi,wo,n,m)
        weight = brdf(wi,wo,n,m,c)*dot_wi_n/pdf
    elif m.type==MaterialType.SPECULAR.value:
        wi = ti.math.reflect(-wo,n)
        weight = ti.math.dot(wi,n)
    elif m.type==MaterialType.BTDF.value:
        wi,n = refraction(wo,ints.n,m.mu)
        weight = m.color
    return weight,Ray(ints.p + n * EPS,wi)

@ti.func
def bxdf_weight(wi:vec3,wo:vec3,ints:Intersection,m:Material,c)->vec3:
    weight = ti.Vector([1.,1.,1.])
    if m.type==MaterialType.BRDF.value:
        weight = brdf(wi,wo,ints.n,m,c)
    elif m.type==MaterialType.SPECULAR.value:
        weight = ti.Vector([0.,0.,0.])
    elif m.type==MaterialType.BTDF.value:
        if ti.math.dot(wo,ints.n)>0:
            weight = brdf(wi,wo,ints.n,m,c) #有反射光泽
        else:
            weight = ti.Vector([0.,0.,0.])

    return weight
